import bm3d
import numpy as np


def bm3d_denoise(image):
    if image.ndim == 3:
        result = np.zeros_like(image)
        for i in range(result.shape[0]):
            result[i] = bm3d.bm3d(image[i], sigma_psd=0.5, stage_arg=bm3d.BM3DStages.ALL_STAGES)
        return result
    elif image.ndim == 2:
        result = bm3d.bm3d(image, sigma_psd=0.5, stage_arg=bm3d.BM3DStages.ALL_STAGES)
        return result
    else:
        raise Exception("Wrong image shape")
